import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from typing import List, Tuple, Dict
import heapq
from flask import Flask, request, jsonify, render_template


class PipelineSystem:
    def __init__(self):
        self.buildings = []
        self.distance_matrix = None
        self.graph = None

    def read_data_from_file(self, filename: str):
        """从文本文件读取建筑物距离矩阵"""
        try:
            with open(filename, 'r', encoding='utf-8') as f:
                lines = f.readlines()

            # 第一行是建筑物名称
            self.buildings = lines[0].strip().split(',')
            n = len(self.buildings)

            # 读取距离矩阵
            self.distance_matrix = np.zeros((n, n))
            for i in range(1, len(lines)):
                row = list(map(float, lines[i].strip().split(',')))
                self.distance_matrix[i - 1] = row

            print("数据读取成功！")
            print(f"建筑物数量: {len(self.buildings)}")

        except Exception as e:
            print(f"读取文件时出错: {e}")

    def read_data_from_text(self, text: str):
        lines = [line.strip() for line in text.splitlines() if line.strip()]
        self.buildings = lines[0].split(',')
        n = len(self.buildings)
        self.distance_matrix = np.zeros((n, n))
        for i in range(1, len(lines)):
            row = list(map(float, lines[i].split(',')))
            self.distance_matrix[i - 1] = row
        self.graph = None

    def build_graph(self):
        """构建图数据结构"""
        if self.distance_matrix is None:
            print("请先读取数据！")
            return

        self.graph = {}
        n = len(self.buildings)

        for i in range(n):
            self.graph[i] = []
            for j in range(n):
                if i != j and self.distance_matrix[i][j] > 0:
                    self.graph[i].append((j, self.distance_matrix[i][j]))

        print("图构建完成！")

    def prim_algorithm(self) -> Tuple[List[Tuple], float]:
        """Prim算法实现"""
        if self.graph is None:
            self.build_graph()

        n = len(self.buildings)
        visited = [False] * n
        min_heap = []
        mst_edges = []
        total_cost = 0

        # 从第一个节点开始
        visited[0] = True
        for neighbor, weight in self.graph[0]:
            heapq.heappush(min_heap, (weight, 0, neighbor))

        print("Prim算法执行过程:")
        step = 1

        while min_heap and len(mst_edges) < n - 1:
            weight, u, v = heapq.heappop(min_heap)

            if not visited[v]:
                visited[v] = True
                mst_edges.append((u, v, weight))
                total_cost += weight

                print(f"步骤 {step}: 连接 {self.buildings[u]} - {self.buildings[v]}, 距离: {weight}")
                step += 1

                # 将新节点的边加入堆
                for neighbor, w in self.graph[v]:
                    if not visited[neighbor]:
                        heapq.heappush(min_heap, (w, v, neighbor))

        return mst_edges, total_cost

    def prim_with_steps(self) -> Tuple[List[Tuple], float, List[Dict]]:
        if self.graph is None:
            self.build_graph()

        n = len(self.buildings)
        visited = [False] * n
        min_heap = []
        mst_edges = []
        total_cost = 0.0
        steps: List[Dict] = []

        start = 0
        visited[start] = True
        for neighbor, weight in self.graph[start]:
            heapq.heappush(min_heap, (weight, start, neighbor))

        while min_heap and len(mst_edges) < n - 1:
            weight, u, v = heapq.heappop(min_heap)
            steps.append({
                "type": "consider",
                "u": u, "v": v, "weight": float(weight),
                "description": f"考虑边 {self.buildings[u]} ←→ {self.buildings[v]}，距离 {weight}"
            })

            if visited[v]:
                steps.append({
                    "type": "reject",
                    "u": u, "v": v, "weight": float(weight),
                    "reason": "目标节点已在生成树中"
                })
                continue

            visited[v] = True
            mst_edges.append((u, v, weight))
            total_cost += weight
            steps.append({
                "type": "select",
                "u": u, "v": v, "weight": float(weight),
                "description": f"选择边 {self.buildings[u]} ←→ {self.buildings[v]}"
            })

            for neighbor, w in self.graph[v]:
                if not visited[neighbor]:
                    heapq.heappush(min_heap, (w, v, neighbor))

        return mst_edges, total_cost, steps

    def kruskal_algorithm(self) -> Tuple[List[Tuple], float]:
        """Kruskal算法实现"""
        if self.graph is None:
            self.build_graph()

        edges = []
        n = len(self.buildings)

        # 收集所有边
        for u in range(n):
            for v, weight in self.graph[u]:
                if u < v:  # 避免重复
                    edges.append((weight, u, v))

        # 按权重排序
        edges.sort()

        # 并查集
        parent = list(range(n))

        def find(x):
            if parent[x] != x:
                parent[x] = find(parent[x])
            return parent[x]

        def union(x, y):
            root_x, root_y = find(x), find(y)
            if root_x != root_y:
                parent[root_y] = root_x
                return True
            return False

        mst_edges = []
        total_cost = 0
        step = 1

        print("Kruskal算法执行过程:")
        for weight, u, v in edges:
            if union(u, v):
                mst_edges.append((u, v, weight))
                total_cost += weight
                print(f"步骤 {step}: 连接 {self.buildings[u]} - {self.buildings[v]}, 距离: {weight}")
                step += 1

                if len(mst_edges) == n - 1:
                    break

        return mst_edges, total_cost

    def kruskal_with_steps(self) -> Tuple[List[Tuple], float, List[Dict]]:
        if self.graph is None:
            self.build_graph()

        edges = []
        n = len(self.buildings)

        for u in range(n):
            for v, weight in self.graph[u]:
                if u < v:
                    edges.append((weight, u, v))
        edges.sort()

        parent = list(range(n))

        def find(x):
            if parent[x] != x:
                parent[x] = find(parent[x])
            return parent[x]

        def union(x, y):
            rx, ry = find(x), find(y)
            if rx != ry:
                parent[ry] = rx
                return True
            return False

        mst_edges = []
        total_cost = 0.0
        steps: List[Dict] = []

        for weight, u, v in edges:
            steps.append({
                "type": "consider",
                "u": u, "v": v, "weight": float(weight),
                "description": f"考虑边 {self.buildings[u]} ←→ {self.buildings[v]}，距离 {weight}"
            })
            if union(u, v):
                mst_edges.append((u, v, weight))
                total_cost += weight
                steps.append({
                    "type": "select",
                    "u": u, "v": v, "weight": float(weight),
                    "description": f"选择边 {self.buildings[u]} ←→ {self.buildings[v]}"
                })
                if len(mst_edges) == n - 1:
                    break
            else:
                steps.append({
                    "type": "reject",
                    "u": u, "v": v, "weight": float(weight),
                    "reason": "形成环路（已在同一连通分量）"
                })

        return mst_edges, total_cost, steps

    def display_solution(self, mst_edges: List[Tuple], total_cost: float, algorithm: str):
        """显示最优方案"""
        print(f"\n=== {algorithm}算法最优方案 ===")
        print(f"总铺设距离: {total_cost:.2f} 米")
        print("铺设路径:")

        for i, (u, v, weight) in enumerate(mst_edges, 1):
            print(f"{i}. {self.buildings[u]} ←→ {self.buildings[v]} : {weight} 米")

    # def plot_solution(self, mst_edges: List[Tuple], algorithm: str):
    #     """绘制最优方案的简易图"""
    #     if not self.buildings:
    #         print("没有数据可绘制！")
    #         return
    #
    #     G = nx.Graph()
    #
    #     # 添加节点
    #     for i, building in enumerate(self.buildings):
    #         G.add_node(i, name=building)
    #
    #     # 添加MST边
    #     for u, v, weight in mst_edges:
    #         G.add_edge(u, v, weight=weight)
    #
    #     # 创建图形
    #     plt.figure(figsize=(12, 8))
    #     pos = nx.spring_layout(G)  # 节点布局
    #
    #     # 绘制节点
    #     nx.draw_networkx_nodes(G, pos, node_size=500, node_color='lightblue')
    #
    #     # 绘制边
    #     nx.draw_networkx_edges(G, pos, width=2, alpha=0.7, edge_color='red')
    #
    #     # 添加标签
    #     labels = {i: self.buildings[i] for i in range(len(self.buildings))}
    #     nx.draw_networkx_labels(G, pos, labels, font_size=8)
    #
    #     # 添加边权重
    #     edge_labels = {(u, v): f'{weight}' for u, v, weight in mst_edges}
    #     nx.draw_networkx_edge_labels(G, pos, edge_labels)
    #
    #     plt.title(f'管线铺设最优方案 - {algorithm}算法', fontsize=15)
    #     plt.axis('off')
    #     plt.tight_layout()
    #     plt.show()

    def plot_solution(self, mst_edges: List[Tuple], algorithm: str):
        """绘制最优方案的简易图"""
        if not self.buildings:
            print("没有数据可绘制！")
            return

        G = nx.Graph()

        # 添加节点
        for i, building in enumerate(self.buildings):
            G.add_node(i, name=building)

        # 添加MST边
        for u, v, weight in mst_edges:
            G.add_edge(u, v, weight=weight)

        # 创建图形
        plt.figure(figsize=(14, 10))

        # 使用预定义的布局以提高一致性
        # 根据数据特点手动设置节点位置，使距离相近的节点在图中也相对靠近
        pos = {
            0: (0, 5),  # 里仁学院教学楼
            1: (2, 5),  # 学生公寓
            2: (4, 5),  # 学生食堂
            3: (6, 6),  # 第一教学楼
            4: (8, 6),  # 第二教学楼
            5: (10, 6),  # 第三教学楼
            6: (12, 6),  # 第四教学楼
            7: (2, 3),  # 材料学院
            8: (4, 3),  # 电气学院
            9: (10, 0),  # 理学院
            10: (8, 0),  # 艺术学院
            11: (6, 0),  # 经管学院
            12: (4, 0),  # 西里西亚学院
            13: (2, 0)  # 信息学院
        }

        # 绘制节点
        nx.draw_networkx_nodes(G, pos, node_size=1200, node_color='lightcoral', alpha=0.9)

        # 绘制边（根据权重调整线宽）
        edges = [(u, v) for u, v, w in mst_edges]
        weights = [w for u, v, w in mst_edges]

        # 归一化权重以适应线宽范围
        max_weight = max(weights) if weights else 1
        min_weight = min(weights) if weights else 0
        width_range = 1 if max_weight == min_weight else max_weight - min_weight

        widths = [2 + 3 * (max_weight - w) / width_range for w in weights]  # 距离越小，线越粗

        nx.draw_networkx_edges(G, pos, edgelist=edges, width=widths, alpha=0.7, edge_color='darkblue')

        # 添加节点标签（建筑物名称）
        labels = {i: self.buildings[i] for i in range(len(self.buildings))}
        nx.draw_networkx_labels(G, pos, labels, font_size=8, font_family='sans-serif')

        # 添加边权重标签
        edge_labels = {(u, v): f'{int(weight)}m' for u, v, weight in mst_edges}
        nx.draw_networkx_edge_labels(G, pos, edge_labels, font_size=7)

        plt.title(f'管线铺设最优方案 - {algorithm}算法', fontsize=16, pad=20)
        plt.axis('off')
        plt.tight_layout()
        plt.margins(0.1)
        plt.show()

    def save_solution_to_file(self, mst_edges: List[Tuple], total_cost: float,
                              algorithm: str, filename: str):
        """将最优方案保存到文件"""
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(f"{algorithm}算法最优管线铺设方案\n")
            f.write(f"总铺设距离: {total_cost:.2f} 米\n")
            f.write("铺设路径详情:\n")

            for i, (u, v, weight) in enumerate(mst_edges, 1):
                f.write(f"{i}. {self.buildings[u]} ←→ {self.buildings[v]} : {weight} 米\n")

        print(f"方案已保存到文件: {filename}")


def create_sample_data_file():
    """创建示例数据文件"""
    sample_data = """里仁学院教学楼,学生公寓,学生食堂,第一教学楼,第二教学楼,第三教学楼,第四教学楼,材料学院,电气学院,理学院,艺术学院,经管学院,西里西亚学院,信息学院
0,350,280,450,520,600,550,800,750,680,720,650,580,620
350,0,150,300,400,450,420,600,580,520,560,500,450,480
280,150,0,200,280,320,300,480,450,400,440,380,350,370
450,300,200,0,120,180,160,350,320,280,310,260,230,250
520,400,280,120,0,90,110,280,250,200,240,190,160,180
600,450,320,180,90,0,80,220,190,150,180,140,120,130
550,420,300,160,110,80,0,180,160,120,150,110,90,100
800,600,480,350,280,220,180,0,60,120,150,200,250,280
750,580,450,320,250,190,160,60,0,90,120,170,220,240
680,520,400,280,200,150,120,120,90,0,80,130,180,200
720,560,440,310,240,180,150,150,120,80,0,100,150,170
650,500,380,260,190,140,110,200,170,130,100,0,80,100
580,450,350,230,160,120,90,250,220,180,150,80,0,70
620,480,370,250,180,130,100,280,240,200,170,100,70,0"""

    with open('campus_distances.csv', 'w', encoding='utf-8') as f:
        f.write(sample_data)
    print("示例数据文件已创建: campus_distances.csv")


def create_app():
    app = Flask(__name__, static_folder='static', template_folder='templates')
    pipeline_system = PipelineSystem()

    @app.route("/")
    def index():
        return render_template("index.html")

    @app.get("/api/sample")
    def sample():
        try:
            pipeline_system.read_data_from_file('campus_distances1.csv')
            pipeline_system.build_graph()
            return jsonify({
                "names": pipeline_system.buildings,
                "matrix": pipeline_system.distance_matrix.tolist()
            })
        except Exception as e:
            return jsonify({"error": str(e)}), 500

    @app.post("/api/compute")
    def compute():
        if request.is_json:
            data = request.get_json()
            names = data.get("names")
            matrix = data.get("matrix")
            algorithm = str(data.get("algorithm", "prim")).lower()

            if not names or not matrix:
                return jsonify({"error": "缺少 names 或 matrix"}), 400

            pipeline_system.buildings = names
            pipeline_system.distance_matrix = np.array(matrix, dtype=float)
            pipeline_system.graph = None
            pipeline_system.build_graph()

        elif "file" in request.files:
            file = request.files["file"]
            text = file.stream.read().decode("utf-8")
            pipeline_system.read_data_from_text(text)
            pipeline_system.build_graph()
            algorithm = str(request.form.get("algorithm", "prim")).lower()
        else:
            return jsonify({"error": "没有提供数据（JSON或上传文件）"}), 400

        if algorithm == "prim":
            mst_edges, total_cost, steps = pipeline_system.prim_with_steps()
            algo_name = "Prim"
        elif algorithm == "kruskal":
            mst_edges, total_cost, steps = pipeline_system.kruskal_with_steps()
            algo_name = "Kruskal"
        else:
            return jsonify({"error": "未知算法"}), 400

        return jsonify({
            "algorithm": algo_name,
            "total_cost": float(total_cost),
            "mst_edges": [{"u": int(u), "v": int(v), "weight": float(w)} for (u, v, w) in mst_edges],
            "steps": steps,
            "buildings": pipeline_system.buildings
        })

    return app


def main():
    """主函数"""
    # # 创建示例数据文件（如果不存在）
    # create_sample_data_file()

    # 创建管线系统实例
    pipeline_system = PipelineSystem()

    # 读取数据
    pipeline_system.read_data_from_file('campus_distances1.csv')

    # 选择算法
    print("\n请选择算法:")
    print("1. Prim算法")
    print("2. Kruskal算法")

    choice = input("请输入选择 (1 或 2): ").strip()

    if choice == '1':
        algorithm_name = "Prim"
        mst_edges, total_cost = pipeline_system.prim_algorithm()
    elif choice == '2':
        algorithm_name = "Kruskal"
        mst_edges, total_cost = pipeline_system.kruskal_algorithm()
    else:
        print("无效选择，使用默认的Prim算法")
        algorithm_name = "Prim"
        mst_edges, total_cost = pipeline_system.prim_algorithm()

    # 显示结果
    pipeline_system.display_solution(mst_edges, total_cost, algorithm_name)

    # 绘制图形
    pipeline_system.plot_solution(mst_edges, algorithm_name)

    # 保存结果到文件
    pipeline_system.save_solution_to_file(mst_edges, total_cost,
                                          algorithm_name, 'pipeline_solution.txt')


if __name__ == "__main__":
    plt.rcParams["font.sans-serif"] = ["SimHei"]  # 设置字体为黑体
    plt.rcParams["axes.unicode_minus"] = False  # 正常显示负号
    app = create_app()
    app.run(host="127.0.0.1", port=5000, debug=True)